import torch
from torch.utils.data import DataLoader, Dataset
import numpy as np
import chatbot.config as config


class ChatbotDataset(Dataset):
    def __init__(self):
        self.input_path = config.chatbot_input_path
        self.target_path = config.chatbot_target_path
        self.input_lines = open(self.input_path, encoding="utf-8").readlines()
        self.target_lines = open(self.target_path, encoding="utf-8").readlines()
        assert len(self.input_lines) == len(self.target_lines), "两者长度不一致"

    def __getitem__(self, index):
        input = self.input_lines[index].strip().split()
        target = self.target_lines[index].strip().split()
        input_length = len(input) if len(input) < config.chatbot_input_max_len else config.chatbot_input_max_len
        target_length = len(target) if len(
            input) < config.chatbot_target_max_len + 1 else config.chatbot_target_max_len + 1
        return input, target, input_length, target_length

    def __len__(self):
        return len(self.input_lines)


def collate_fn(batch):
    batch = sorted(batch, key=lambda x: x[-2], reverse=True)
    input, target, input_length, target_length = zip(*batch)
    input = [config.chatbot_ws_input.transform(i, max_len=config.chatbot_input_max_len) for i in input]
    target = [config.chatbot_ws_target.transform(i, max_len=config.chatbot_target_max_len, add_eos=True) for i in
              target]
    return torch.LongTensor(input).to(config.device), torch.LongTensor(target).to(config.device), torch.LongTensor(
        input_length).to(config.device), torch.LongTensor(
        target_length).to(config.device)


train_data_loader = DataLoader(ChatbotDataset(), batch_size=config.chatbot_batch_size, shuffle=True,
                               collate_fn=collate_fn)

if __name__ == '__main__':
    # ds = NumDataset()
    # print(ds.data[:10])
    # print(ds[0])
    # print(len(ds))
    for input, label, input_length, label_length in train_data_loader:
        print(input)
        print(label)
        print(input_length)
        print(label_length)
        break
